import torch
import reversi_layer_cpp


@torch.no_grad()
def evaluate(model, test_size):
    model.eval()
    data_b = torch.zeros((test_size,), dtype=torch.long)
    data_w = torch.zeros((test_size,), dtype=torch.long)
    data_b[:] = 0x0000000810000000
    data_w[:] = 0x0000001008000000
    reversi_layer_cpp.std(data_b[:test_size//2], data_w[:test_size//2])
    is_std = True
    win_cnt = 0.
    lose_cnt = 0.
    while data_b.size(0):
        if is_std:
            win = reversi_layer_cpp.std(data_b, data_w)
            lose_cnt += torch.sum(win == 2)
            win_cnt += torch.sum(win == 0)
        else:
            pred = model(reversi_layer_cpp.convert(data_b, data_w).half().to('cuda')).to('cpu').float()
            _, win, _ = reversi_layer_cpp.forward(data_b, data_w, pred)
            win_cnt += torch.sum(win == 2)
            lose_cnt += torch.sum(win == 0)
        data_b = data_b[win < 0]
        data_w = data_w[win < 0]
        is_std = not is_std

    win = reversi_layer_cpp.win(data_b, data_w)
    model.train()
    return (test_size/2 - lose_cnt/2 + win_cnt/2) / test_size
